import os
import cv2
import numpy as np
import torch
import torch.utils.data as torch_data
import torchvision.transforms as tf
import PIL.Image
import PIL.ImageOps
import matplotlib.pyplot as plt
import json
import utils.datasets_utils as data_utils


class BarcodeKeypointDataset(torch_data.Dataset):
    """用于生成条码关键点检测数据集，其升级网络的输出为(B,C=12,H,W),三类条码每一类对应2张位置热图+2张方向热图"""

    def __init__(
        self,
        dataset_dir: str,
        transform=None,
        input_size=(128, 128),
        output_size=(128, 128),
        sigma=4.0,
    ):
        """
        图片关键点数据类初始化函数
        :param dataset_dir: 包含图片与标签的文件夹路径
        :param transform: 图片转换为Tensor矩阵的流程
        :param input_size: 神经网络输出图片的尺寸(宽、高)
        :param output_size: 神经网络输出热图的尺寸(宽、高)
        """
        super(BarcodeKeypointDataset, self).__init__()
        self.dataset_dir = dataset_dir  #
        self.json_files = [
            f for f in os.listdir(dataset_dir) if f.endswith(".json")
        ]  # json 标签名称
        self.transform = transform  # torchvision.transforms中选取转换操作
        self.input_size = input_size  # 输入神经网络图片的尺寸
        self.output_size = output_size  # 输出heatmap的尺寸
        self.sigma = (
            sigma  # 高斯核的标准差，决定热力图的大小/扩散，随训练轮数增加而减小
        )

    def __len__(self):
        """返回训练数据图片数量"""
        return len(self.json_files)

    def __getitem__(self, idx):
        """
        构造指定索引样本的模型输入。
        :param idx: 样本的索引。
        :return: 返回一个元组，包含以下元素：
            input_image (tensor): 预处理后输入神经网络的图像
            img_heatmap (ndarray): 生成的热图数组，包含了所有码制的位置热图和方向热图。
        """
        # label_path对应的json文件，为符合labelme标注格式的字典,详见
        # https://doc.weixin.qq.com/doc/w3_AQ0ArAaVAOcT8KTg7BmSxazxA8e6t?scode=AAsAKQcQABILJY68yh
        label_path = os.path.join(self.dataset_dir, self.json_files[idx])
        # 读取json文件并转化为字典labels
        with open(label_path, "r") as f:
            barcodes_data_labels = json.load(f)
        image_path = os.path.join(
            self.dataset_dir, barcodes_data_labels["imagePath"]
        )
        # 读取图像并转换为灰度图
        image = PIL.Image.open(
            image_path,
        ).convert("L")

        # 原始图像尺寸
        original_width, original_height = image.size
        # 输入神经网络的图像尺寸
        input_width, input_height = self.input_size
        # 输出神经网络的heatmap尺寸
        output_width, output_height = self.output_size

        if input_width / original_width != input_height / original_height:
            raise ValueError("原始图像与神经网络输入图像高宽比不一致")
        if output_width / original_width != output_height / original_height:
            raise ValueError(
                "神经网络输出heatmaps高宽比需要与网络输入图像高宽比一致"
            )

        # 神经网络输入尺寸与原始图像尺寸的缩放比例
        scale = input_width / original_width
        # 使用双线性插值将图像等比例缩放为神经网络输入尺寸
        input_image = image.resize(
            (round(original_width * scale), round(original_height * scale)),
            PIL.Image.Resampling.BILINEAR,
        )

        # 将图片转为Tensor
        if self.transform:
            input_image = self.transform(input_image)

        # 分别生成各个条码的heatmap和权重值
        # 一维条码的热图初始化
        (
            heatmap_code_1D_13,
            heatmap_code_1D_24,
            heatmap_code_1D_max,
            heatmap_code_1D_sin,
            heatmap_code_1D_cos,
        ) = data_utils.initialize_heatmaps(self.output_size)
        # 建立keypoint_id与heatmap的映射
        keypoint_id_to_heatmap_1D = {
            0: heatmap_code_1D_13,
            1: heatmap_code_1D_24,
            2: heatmap_code_1D_13,
            3: heatmap_code_1D_24,
        }

        # 从字典中读取存有关键点信息的值，”shapes“为对应的键
        barcodes_data = barcodes_data_labels["shapes"]

        # 神经网络输出尺寸与原始图像尺寸的缩放比例
        # 用于调整标签中关键点的坐标参数，并生成heatmap的ground truth
        scale_original_to_output = output_width / original_width

        # barcodes_data为一列表，列表元素为一字典，字典中的”label“键对应条码名称，
        # ”points“键对应了关键点列表,先观察标签json文件可以理解。
        # 依次遍历列表每一个元素，得到最终的热图
        for barcode_data in barcodes_data:
            # 获取原图中关键点坐标，并对其坐标进行缩放，对应到输出热图中
            keypoint_coords = barcode_data["points"]
            keypoint_coords = (
                np.array(keypoint_coords) * scale_original_to_output
            )

            # 根据关键点坐标求x,y坐标均值，得到条码的中心点坐标
            center_coords = np.mean(keypoint_coords, axis=0)

            if barcode_data["label"] in [
                "CODE128",
                "CODE39",
                "EAN13_GUARD",
                "EAN8_GUARD",
            ]:
                # 遍历4个关键点，更新关键点对应的位置热图和方向热图。
                for (
                    keypoint_id,
                    heatmap_part,
                ) in keypoint_id_to_heatmap_1D.items():
                    (
                        heatmap_part,
                        heatmap_code_1D_max,
                        heatmap_code_1D_sin,
                        heatmap_code_1D_cos,
                    ) = data_utils.update_heatmaps_and_direction(
                        keypoint_coords[keypoint_id],
                        center_coords,
                        heatmap_part,
                        heatmap_code_1D_max,
                        heatmap_code_1D_sin,
                        heatmap_code_1D_cos,
                        sigma=self.sigma,
                    )
            elif barcode_data["label"] == "QR":
                pass
            elif barcode_data["label"] == "DM":
                pass
            elif barcode_data["label"] == "WHITE_LABEL":
                pass
            else:
                raise ValueError("不支持的条码类型")

        image_heatmap_1234 = np.stack(
            [
                heatmap_code_1D_13,  # 一维条码位置热图part1
                heatmap_code_1D_24,  # 一维条码位置热图part2
            ]
        )
        image_heatmap_1234 = np.max(image_heatmap_1234, axis=0)

        image_heatmap_sin = np.stack([heatmap_code_1D_sin]).squeeze()

        image_heatmap_cos = np.stack([heatmap_code_1D_cos]).squeeze()
        img_heatmap = np.stack(
            [
                image_heatmap_1234,  # 一维条码位置热图part1
                image_heatmap_sin,  # 一维条码位置热图part2
                image_heatmap_cos,
            ],
            axis=0,
        )
        return input_image, img_heatmap


def visualize_heatmaps(image, heatmaps, score_threshold):
    """
    可视化缩放后的原图，并在其上绘制条码关键点位置热图，画箭头可视化关键点指向的方向。
    （目前仅绘制了QR码）
    :param image: 输入图像 (H, W, C)
    :param heatmaps: 关键点热力图 (3, H, W)，3个通道代表了1个位置热力图和2个角度热力图
    :param score_threshold: 热力图中的需要绘制箭头的点的阈值，默认为1，一个关键点画一条线
    """

    # 缩放原图以匹配热力图的大小
    heatmap_height, heatmap_width = heatmaps[0].shape
    resized_image = cv2.resize(image, (heatmap_width, heatmap_height))

    # 在原图上绘制所有关键点的方向
    for i in range(0, 2):  # 绘制qr的方向
        y_coords, x_coords = np.where(heatmaps[i] >= score_threshold)
        for y, x in zip(y_coords, x_coords):
            sin_val = heatmaps[1][y, x]  # QR的sin值
            cos_val = heatmaps[2][y, x]  # QR的cos值
            angle_rad = np.arctan2(sin_val, cos_val)
            # 计算箭头的终点坐标，10为箭头的长度
            dx = x.item() + np.cos(angle_rad) * 10
            dy = y.item() + np.sin(angle_rad) * 10
            cv2.arrowedLine(
                resized_image, (x, y), (int(dx), int(dy)), (0, 0, 255), 2
            )
    # 展示原图以及QR码对应的位置热图与方向热图
    plt.figure("-1")
    plt.imshow(resized_image, cmap="gray")
    plt.figure("1d_heatmap")
    plt.imshow(heatmaps[0])
    plt.figure("1d_sin")
    plt.imshow(heatmaps[1])
    plt.figure("1d_cos")
    plt.imshow(heatmaps[2])
    plt.show()


def collate_fn(batch):
    """
    将一个batch的数据转换为模型的输入。
    :param batch：一个batch的数据，是一个列表，每个元素是一个元组，包含了图像、关键点热图。

    :return: 转换后的数据，包含了图像、关键点热图。
    """
    images, heatmaps = zip(*batch)

    images = [image.clone().detach() for image in images]
    heatmaps = [torch.tensor(heatmap) for heatmap in heatmaps]

    return torch.stack(images), torch.stack(heatmaps)


# 用于测试的函数
def t_dataset():
    # 将图片转为(C*H*W)的tensor
    transform = tf.ToTensor()
    dataset = BarcodeKeypointDataset(
        r"..\barcode_val_datasets", transform=transform, sigma=4
    )
    dataloader = torch_data.DataLoader(
        dataset, batch_size=1, shuffle=True, collate_fn=data_utils.collate_fn
    )
    for batch_images, batch_heatmaps in dataloader:
        for image, heatmaps in zip(batch_images, batch_heatmaps):
            image_np = image.numpy().transpose(
                1, 2, 0
            )  # 将C, H, W格式转为H, W, C
            heatmaps_np = heatmaps.numpy()
            visualize_heatmaps(image_np, heatmaps_np, 1)
        break


if __name__ == "__main__":
    t_dataset()
